from pynput import keyboard
from interbotix_xs_modules.xs_robot.turret import InterbotixTurretXS

def on_press(key):
    global pan_position, tilt_position

    try:
        if key.char == 'w':
            tilt_position = max(tilt_position - step, tilt_limits[0])
            robot.turret.tilt(position=tilt_position, blocking=False)

        elif key.char == 's':
            tilt_position = min(tilt_position + step, tilt_limits[1])
            robot.turret.tilt(position=tilt_position, blocking=False)

        elif key.char == 'a':
            pan_position = min(pan_position + step, pan_limits[1])
            robot.turret.pan(position=pan_position, blocking=False)

        elif key.char == 'd':
            pan_position = max(pan_position - step, pan_limits[0])
            robot.turret.pan(position=pan_position, blocking=False)

    except AttributeError:
        pass

def on_release(key):
    if key == keyboard.Key.esc or (hasattr(key, 'char') and key.char == 'q'):
        # Stop listener
        return False

def main():
    global robot, pan_position, tilt_position, step, pan_limits, tilt_limits

    robot = InterbotixTurretXS(
        robot_model="pxxls",
        pan_profile_type="time",
        tilt_profile_type="time"
    )

    pan_position = 0.0
    tilt_position = 0.0
    step = 0.3

    pan_limits = (-3.14, 3.14)
    tilt_limits = (-1.57, 1.57)

    print("按键控制：w/a/s/d移动，q退出。按Esc终止监听。")

    # Collect events until released
    with keyboard.Listener(
            on_press=on_press,
            on_release=on_release) as listener:
        listener.join()

    robot.shutdown()

if __name__ == '__main__':
    main()

